#!/usr/bin/env python
# coding: utf-8

# 数值方法3：偏微分方程1 使用有限差分法解一维热传导（扩散）方程
# 无热源情况

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation
from matplotlib.animation import FFMpegWriter
h = 8.6125 # 空气交换系数
# 录入材料信息
# 根据厚度对介质进行分割，方便x轴标数计算
m = np.array([6, 60, 36, 50])
# 分割后计算总厚度，用于规定X的范围
m_sum = np.sum(m)
# 对时间进行分割和规定总时长
n = 3600
time = 3600
# 材料的厚度
# 四种材料厚度
l1 = 0.6/1000
l2 = 6/1000
l3 = 3.6/1000
l4 = 5/1000
l = np.array([l1, l2, l3, l4])
# 四种材料的热传导率
lam_1 = 0.082
lam_2 = 0.37
lam_3 = 0.045
lam_4 = 0.028
lam = np.array([lam_1, lam_2, lam_3, lam_4])
# 四种材料的密度
de_1 = 300
de_2 = 862
de_3 = 74.2
de_4 = 1.18
de = np.array([de_1, de_2, de_3, de_4])
# 四种材料的比热容
c1 = 1377
c2 = 2100
c3 = 1726
c4 = 1005
c = np.array([c1, c2, c3, c4])


# 计算热扩率 && 计算材料长度分割和时间步长分割求解 && 计算各层介质剖分的步长比
a = []  # 热扩率
derta_x = []  # 材料分割步长
derta_t = time/n  # 时间分割步长
r = []  # 介质剖分步长
u = np.zeros([m_sum, n])  # 定义四层耦合介质温度分布矩阵 每一列是一个x轴
for i in range(len(c)):

    temp_a = lam[i]/(c[i]*de[i])
    temp_x = l[i] / m[i]
    temp_r = derta_t/np.square(temp_x)*temp_a

    a.append(temp_a)
    derta_x.append(temp_x)
    r.append(temp_r)
a = np.array(a)
derta_x = np.array(derta_x)
r = np.array(r)
# from IPython import get_ipython
'''
1、初始热量
2、边界条件
3、无热源
4、计算矩阵A
'''
# 根据公式推导出矩阵A，用于计算时间变化的热量
'''A = (-2 * np.eye(len(x), len(x), dtype=int)
     + np.diag(np.diag(np.ones([len(x), len(x)]), k=1), k=1)
     + np.diag(np.diag(np.ones([len(x), len(x)]), k=1), k=-1))'''

A = np.zeros([m_sum, m_sum])
# 第一个边界
for i in range(0, m[0]-1):
    A[i, i] = 1+2*r[0]
    A[i, i+1] = -r[0]
    if i >= 1:
        A[i, i-1] = -r[0]
# 第一个边界的边界条件
A[m[0]-1, m[0]-1
] = (lam[0]/derta_x[0]+lam[1]/derta_x[1])
A[m[0]-1, m[0]-2] = lam[0]/derta_x[0]
A[m[0]-1, m[0]] = lam[1]/derta_x[1]

# 第二个边界
for i in range(m[0], m[1]+m[0]-1):
    A[i, i] = 1 + 2 * r[1]
    A[i, i + 1] = -r[1]
    A[i, i - 1] = -r[1]
# 第二个边界的边界条件
A[m[1]+m[0]-1, m[1]+m[0]-1] = (lam[2]/derta_x[2]+lam[1]/derta_x[1])
A[m[1]+m[0]-1, m[1]+m[0]-2] = lam[1]/derta_x[1]
A[m[1]+m[0]-1, m[1]+m[0]] = lam[2]/derta_x[2]

# 第三个边界
for i in range(m[1]+m[0], m[1]+m[0]+m[2]-1):
    A[i, i] = 1 + 2 * r[2]
    A[i, i + 1] = -r[2]
    A[i, i - 1] = -r[2]
# 第三个边界的边界条件
A[m[2]+m[1]+m[0]-1, m[2]+m[1]+m[0]-1] = (lam[3]/derta_x[3]+lam[2]/derta_x[2])
A[m[2]+m[1]+m[0]-1, m[2]+m[1]+m[0]-2] = lam[2]/derta_x[2]
A[m[2]+m[1]+m[0]-1, m[2]+m[1]+m[0]] = lam[3]/derta_x[3]

# 第四个边界
for i in range(m[2]+m[1]+m[0], m[3]+m[1]+m[0]+m[2]-1):
    A[i, i] = 1 + 2 * r[3]
    A[i, i + 1] = -r[3]
    A[i, i - 1] = -r[3]
A[m_sum-1, m_sum-1] = h + lam[3]/derta_x[3]
A[m_sum-1, m_sum-2] = lam[3]/derta_x[3]

# 解方程组
u = np.mat(np.ones([m_sum, 1])*35)  # 将初热值转置为列向量
u[0, 0] = 75
# print(u.shape)
for n in range(0,time-1):
    print(np.mat(A).shape,u[:,n].shape)
    u = np.append(u, u[:, n] + np.mat(A) * u[:, n]*(a * a * derta_t / 0.0001 / 0.0001), axis=1)  # 根据传导公式解除随时间变化的向量
    # 对边界值进行特殊处理
    u[0, n + 1] = 75

# 绘制边界线
Y = np.array(u)  # 计算后得到的温度值
X = np.array(range(0,m_sum))  # 位置
plt.plot(X, Y[:, -1])  # 提取的是最终结果的温度
plt.title('FinalStates Temperature')
plt.savefig("FinalStatesT.png")  # 保存图片
plt.show()

# 3D曲面
# get_ipython().run_line_magic('matplotlib', 'notebook')
# 准备数据
t = np.array(range(0,time))
X_x = np.array(range(0,m_sum))
x_x, y_y = np.meshgrid(t, X_x)

# 绘制图片
fig = plt.figure("3D Surface", facecolor="lightgray")
plt.title("Temperature3D", fontsize=18)

# 设置为3D图片类型
ax3d = Axes3D(fig, auto_add_to_figure=False)
fig.add_axes(ax3d)

ax3d.set_xlabel("X")
ax3d.set_ylabel("Y")
ax3d.set_zlabel("Z")
plt.tick_params(labelsize=10)

ax3d.plot_surface(x_x, y_y, u, cstride=20, rstride=20, cmap="jet")

# plt.savefig("Temperature3D.png")  # 保存图片
plt.show()

# 初始化画布
fig = plt.figure()
# plt.xlim(0,X)
plt.ylim(0, np.max(u))  # np.max(psi)
plt.grid(ls='--')

ys = u[:, 0]
Figure = plt.plot(x, ys, c='blue', alpha=0.8)[0]


# 更新函数
def updata(num):
    ys = u[:, num]
    Figure.set_data(x, ys)
    return Figure


# animation库绘制动图
ani = animation.FuncAnimation(fig=fig, func=updata, frames=np.arange(0, int(len(t) / 2)),
                              interval=5)  # frames帧数，interval间隔（帧）
plt.show()

## 保存为mp4，运行速度较慢，不保存时注释掉。
# mywriter = FFMpegWriter(fps=60)
# ani.save('TemperatureHistory.MP4',writer=mywriter)